from PySide2.QtWidgets import QApplication, QMessageBox
from PySide2.QtUiTools import QUiLoader
from PySide2.QtWidgets import QFileDialog
import os
import subprocess
import cv2 as cv
import re
from PyQt5.QtWidgets import (QWidget, QApplication, QColorDialog, QPushButton, QFrame)
from PyQt5.QtGui import QColor

# 时间:20220927
# 功能实现: 选择源图像位置、目的图像位置，输入上下左右边框像素值，输入RGB值，可以为图像添加边框效果
#         若输出位置为空，则默认输出到源图像位置，且文件名后为_tran
#         多级目录图像的识别

#20220928：添加了程序异常的提示信息；添加了可自定义图像格式

# 20220930：添加了图像处理输出功能；修复边框参数/RGB参数输入校验；

# 20221013:修复边框参数/RGB参数输入校验；添加颜色预览

class Stats:
    def __init__(self):
        self.ui = QUiLoader().load('主界面UI.ui')
        self.ui.src_dir_button.clicked.connect(self.srcDirChoose)
        self.ui.dst_dir_button.clicked.connect(self.dstDirChoose)
        self.ui.start.clicked.connect(self.start)
        self.ui.del_all.clicked.connect(self.delAll)
        self.ui.clear_log.clicked.connect(self.clearLog)
        # 上下左右边框值的校验
        borders = [self.ui.top,self.ui.down,self.ui.left,self.ui.right]
        for border in borders:
            border.textChanged.connect(self.borders_check)

        color_rgb = [self.ui.R_color,self.ui.G_color,self.ui.B_color]
        for color in color_rgb:
            # RGB文本值正确性校验
            color.textChanged.connect(self.rgb_color_number_check)
            # RGB对应的预览颜色
            color.textChanged.connect(self.color_setup)

        # 颜色板默认颜色
        col = QColor(255,255,255)
        self.ui.select_color.clicked.connect(self.showDailog)
        self.ui.color.setStyleSheet("QWidget { background-color: %s }" % col.name())

    def borders_check(self):
        for border in [self.ui.top,self.ui.down,self.ui.left,self.ui.right]:
            check_border = re.search('^[+]{0,1}(\d+)$',border.text())
            if check_border is None:
                QMessageBox.information(
                    self.ui,
                    '错误',
                    '边框值必须为正整数')
                border.setText("10")
                return

    def rgb_color_number_check(self):
        # 检查RGB数值
        for rgb_color in [self.ui.B_color,self.ui.G_color,self.ui.R_color]:
            check_color = re.search('^[+]{0,1}(\d+)$',rgb_color.text())
            if check_color is None:
                QMessageBox.information(
                    self.ui,
                    '错误',
                    'RGB值必须为正整数，范围为[0~255]')
                return
            if int(rgb_color.text()) >255:
                rgb_color.setText('255')

    def color_setup(self):
        hex_color_temp = []
        b_color = int(self.ui.B_color.text())
        g_color = int(self.ui.G_color.text())
        r_color = int(self.ui.R_color.text())
        color = [r_color,g_color,b_color]

        for i in color:
            i = hex(int(i))
            if '0x' in i:
                i = i.replace('0x', '')
            hex_color_temp.append(str(i).zfill(2))
        hex_color = ''.join(hex_color_temp)
        self.ui.color.setStyleSheet("QWidget { background-color: #%s }" % hex_color)

    # 源图像路径选择
    def srcDirChoose(self):
        filePath = QFileDialog.getExistingDirectory(self.ui, "选择源图像路径")
        self.ui.src_dir.setText(filePath)

    # 目的图像路径选择
    def dstDirChoose(self):
        filePath = QFileDialog.getExistingDirectory(self.ui, "选择目的图像路径")
        self.ui.dst_dir.setText(filePath)

    # 颜色板
    def showDailog(self):
        '''弹出一个ColorDailog 对话框，把获取的颜色，给clo'''
        col = QColorDialog.getColor(title='取色面板')

        if col.isValid():  # 当clo获取到了有效的颜色
            # 设置self.frm的颜色为clo获取到的颜色
            self.ui.color.setStyleSheet("QWidget { background-color: %s }" % col.name())
            # self.color_transfer(col.name())
        self.ui.R_color.setText(str(int(col.name()[1:3],16)))
        self.ui.G_color.setText(str(int(col.name()[3:5],16)))
        self.ui.B_color.setText(str(int(col.name()[5:7],16)))

    # 转换程序
    def start(self):
        # 检查源图像路径 则报提示
        if self.ui.src_dir.text()=='':
            QMessageBox.critical(
                self.ui,
                '错误',
                '源图像位置不允许为空！')
            return

        top = int(self.ui.top.text())
        down = int(self.ui.down.text())
        left = int(self.ui.left.text())
        right = int(self.ui.right.text())

        b_color = int(self.ui.B_color.text())
        g_color = int(self.ui.G_color.text())
        r_color = int(self.ui.R_color.text())
        color = [b_color, g_color, r_color]

        # 读取目录所有文件
        imgs = []
        # 文件类型
        file_type = ['jpg','JPG','png','PNG']
        add_file_type = self.ui.filetype.text()
        file_type = file_type + add_file_type.split(' ')

        for dir_name, sub_name, files_name in os.walk(self.ui.src_dir.text()):
            dir_name = dir_name.replace('\\','/')
            if self.ui.dst_dir.text()!='':
                # 如果源路径中有多层目录，则在输出目录中创建多层目录结构
                if not os.path.exists(self.ui.dst_dir.text()+dir_name.split(self.ui.src_dir.text())[1]):
                    os.mkdir(self.ui.dst_dir.text()+dir_name.split(self.ui.src_dir.text())[1])
            # 返回当前目录，子目录，文件名
            for file_name in files_name:
                if file_name.split('.')[-1] in file_type:
                    imgs.append(dir_name + '/'+file_name)

        if self.ui.dst_dir.text()!='':
            img_dir = self.ui.dst_dir.text()
        else:
            img_dir = self.ui.src_dir.text()

        # 对图像进行处理
        for img in imgs:
            img1 = cv.imread(img)
            constant = cv.copyMakeBorder(img1, top, down, left, right, cv.BORDER_CONSTANT, value=color)

            # 图像处理及保存
            print(img_dir + img[img.rfind("/"):img.rfind(".")]+'_tran'+img[img.rfind("."):])
            print(img[img.rfind("."):])
            try:
                cv.imwrite(img_dir + img[img.rfind("/"):img.rfind(".")]+'_tran'+img[img.rfind("."):], constant)
                self.ui.loginfo.appendPlainText(
                    '已输出到：' + img_dir + img[img.rfind("/"):img.rfind(".")] + '_tran' + img[img.rfind(
                        "."):])
            except:
                cv.imwrite(img_dir + img[img.rfind("/"):img.rfind(".")]+'_tran'+'.png', constant)
                self.ui.loginfo.appendPlainText(
                    '已输出到：' + img_dir + img[img.rfind("/"):img.rfind(".")] + '_tran.png')



        check_val = self.check_img_save(imgs)
        if check_val == 1:
            choice = QMessageBox.question(self.ui, '提示', '已转换完成')
            if choice == QMessageBox.Yes:
                if self.ui.dst_dir.text() != '':
                    os.startfile(self.ui.dst_dir.text())
                else:
                    os.startfile(self.ui.src_dir.text())
        else:
            QMessageBox.critical(self.ui, "错误", "可能是目录或文件名包含中文")


    # 删除所有转换后的图像
    def delAll(self):
        # 判断目的图像路径是否为空
        if self.ui.dst_dir.text()!='':
            dir = self.ui.dst_dir.text()
        else:
            dir = self.ui.src_dir.text()

        for dir_name, sub_name, files_name in os.walk(dir):
            dir_name = dir_name.replace('\\', '/')
            # 返回当前目录，子目录，文件名
            for file_name in files_name:
                if '_tran' in file_name:
                    del_img_names = os.path.join(dir_name, file_name)
                    del_img_names = del_img_names.replace('\\', '/')
                    os.remove(del_img_names)

    # 清空输日志信息
    def clearLog(self):
        self.ui.loginfo.clear()

    # 检查文件是否存在
    def check_img_save(self,imgs):
        check_val = 1
        if self.ui.dst_dir.text() == '':
            dir = self.ui.src_dir.text()
        else:
            dir = self.ui.dst_dir.text()

        for img in imgs:
            # 判断在指定输出位置 能否找到输出的图像
            # print(dir + img[img.rfind("/"):img.rfind(".")] + '_tran' + img[img.rfind("."):])
            if os.path.exists(dir + img[img.rfind("/"):img.rfind(".")] + '_tran' + img[img.rfind("."):]):
                check_val = check_val * 1
            else:
                check_val = check_val * 0
                return check_val
            return check_val


app = QApplication([])
stats = Stats()
stats.ui.show()
app.exec_()

